/*
 * MIT License
 *
 * Copyright (c) 2021 cao.yong<cymsn81@126.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package org.cy.websocket.stomp.messaging;

import org.cy.websocket.stomp.StompCommand;
import org.cy.websocket.stomp.StompHeaders;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.UnsupportedEncodingException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;

/**
 * Created in 2021-05-09 21:39
 *
 * @author cao.yong
 */
public class StompDecoder {
    private static final Logger LOG = LoggerFactory.getLogger(StompDecoder.class);

    static final byte[] HEARTBEAT_PAYLOAD = new byte[]{'\n'};

    public Message<byte[]> decode(byte[] bytes) {
        Message<byte[]> message = null;
        ByteBuffer byteBuffer = ByteBuffer.wrap(bytes);
        skipEol(byteBuffer);

        byteBuffer.mark();

        String command = readCommand(byteBuffer);
        if (command.length() > 0) {
            StompCommand stompCommand = StompCommand.valueOf(command);
            StompHeaders headers = new StompHeaders();
            byte[] payload = null;
            if (byteBuffer.remaining() > 0) {
                readHeaders(byteBuffer, headers);
                payload = readPayload(byteBuffer, headers);
            }
            if (payload != null) {
                message = new GenericMessage<>(stompCommand, headers, payload);
            } else {
                byteBuffer.reset();
            }
        } else {
            message = new GenericMessage<>(StompCommand.HEARTBEAT, HEARTBEAT_PAYLOAD);
        }
        return message;
    }

    protected void skipEol(ByteBuffer byteBuffer) {
        while (true) {
            if (!tryConsumeEndOfLine(byteBuffer)) {
                break;
            }
        }
    }

    private String readCommand(ByteBuffer byteBuffer) {
        ByteArrayOutputStream command = new ByteArrayOutputStream(256);
        while (byteBuffer.remaining() > 0 && !tryConsumeEndOfLine(byteBuffer)) {
            command.write(byteBuffer.get());
        }
        try {
            return command.toString(StandardCharsets.UTF_8.name());
        } catch (UnsupportedEncodingException e) {
            throw new IllegalArgumentException("Invalid charset name: " + StandardCharsets.UTF_8.name(), e);
        }
    }

    private void readHeaders(ByteBuffer byteBuffer, StompHeaders headers) {
        while (true) {
            ByteArrayOutputStream headerStream = new ByteArrayOutputStream(256);
            boolean headerComplete = false;
            while (byteBuffer.hasRemaining()) {
                if (tryConsumeEndOfLine(byteBuffer)) {
                    headerComplete = true;
                    break;
                }
                headerStream.write(byteBuffer.get());
            }
            if (headerStream.size() > 0 && headerComplete) {
                try {
                    String header = headerStream.toString(StandardCharsets.UTF_8.name());
                    int colonIndex = header.indexOf(':');
                    if (colonIndex <= 0) {
                        if (byteBuffer.remaining() > 0) {
                            throw new StompConversionException("Illegal header: '" + header + "'. A header must be of the form <name>:[<value>].");
                        }
                    } else {
                        String headerName = unescape(header.substring(0, colonIndex));
                        String headerValue = unescape(header.substring(colonIndex + 1));
                        headers.add(headerName, headerValue);
                    }
                } catch (UnsupportedEncodingException e) {
                    throw new StompConversionException(e.getMessage(), e);
                }
            } else {
                break;
            }
        }
    }

    private byte[] readPayload(ByteBuffer byteBuffer, StompHeaders headers) {
        Integer contentLength;
        try {
            contentLength = headers.getContentLength();
        } catch (NumberFormatException ex) {
            if (LOG.isDebugEnabled()) {
                LOG.debug("Ignoring invalid content-length: '" + headers);
            }
            contentLength = null;
        }
        if (null != contentLength && contentLength >= 0) {
            if (byteBuffer.remaining() > contentLength) {
                byte[] payload = new byte[contentLength];
                byteBuffer.get(payload);
                if (byteBuffer.get() != 0) {
                    throw new StompConversionException("Frame must be terminated with a null octet");
                }
                return payload;
            }
        } else {
            ByteArrayOutputStream payload = new ByteArrayOutputStream(256);
            while (byteBuffer.remaining() > 0) {
                byte b = byteBuffer.get();
                if (b == 0) {
                    return payload.toByteArray();
                } else {
                    payload.write(b);
                }
            }
        }
        return null;
    }

    private boolean tryConsumeEndOfLine(ByteBuffer byteBuffer) {
        if (byteBuffer.remaining() > 0) {
            byte b = byteBuffer.get();
            if (b == '\n') {
                return true;
            } else if (b == '\r') {
                if (byteBuffer.remaining() > 0 && byteBuffer.get() == '\n') {
                    return true;
                } else {
                    throw new StompConversionException("'\\r' must be followed by '\\n'");
                }
            }
            // Explicit cast for compatibility with covariant return type on JDK 9's ByteBuffer
            byteBuffer.position(byteBuffer.position() - 1);
        }
        return false;
    }

    private String unescape(String inString) {
        StringBuilder sb = new StringBuilder(inString.length());
        int pos = 0;  // position in the old string
        int index = inString.indexOf('\\');

        while (index >= 0) {
            sb.append(inString, pos, index);
            if (index + 1 >= inString.length()) {
                throw new StompConversionException("Illegal escape sequence at index " + index + ": " + inString);
            }
            char c = inString.charAt(index + 1);
            if (c == 'r') {
                sb.append('\r');
            } else if (c == 'n') {
                sb.append('\n');
            } else if (c == 'c') {
                sb.append(':');
            } else if (c == '\\') {
                sb.append('\\');
            } else {
                // should never happen
                throw new StompConversionException("Illegal escape sequence at index " + index + ": " + inString);
            }
            pos = index + 2;
            index = inString.indexOf('\\', pos);
        }

        sb.append(inString.substring(pos));
        return sb.toString();
    }
}
